from os.path import join, abspath
import json
from typing import List
from backend.experiment.framework.shots_prompt.task import \
    PerplexityPromptShotsChoiceMakingTaskMaker
from backend.experiment.framework.task import TaskMaker
from backend.experiment.tokenizer import tokenizer
from backend.experiment.framework.shots_prompt.implement import \
    CommonPerplexityChoiceMakingImplement, \
    CommonShotsGenerator, \
    CommonPerplexityChoiceMakingSample
from backend.experiment.model import context_limit


class C3Implement(CommonPerplexityChoiceMakingImplement):
    __mask_token = '[MASK]'

    @classmethod
    def _make_test_sample(cls, record, made_sample_num: int) -> \
            List[CommonPerplexityChoiceMakingSample]:
        res = []
        text = '\n'.join(record[0])
        questions = record[1]
        for each in questions:
            question = each['question']
            choice = each['choice']
            prompt = cls.__make_single_prompt(text, question, cls.__mask_token)
            assert len(prompt) < 1024, f'prompt: {prompt}\n len: {len(prompt)}'
            res.append(
                CommonPerplexityChoiceMakingSample(
                    prompt=prompt,
                    target_mask_token=cls.__mask_token,
                    sample_id=made_sample_num,
                    candidates=choice
                )
            )
            made_sample_num += 1
        return res

    @classmethod
    def _read_train_records(cls) -> list:
        with open(join(cls.data_dir(), 'm-train.json'), 'r', encoding='utf8') \
                as file:
            records = json.loads(file.read())
        with open(join(cls.data_dir(), 'd-train.json'), 'r', encoding='utf8') \
                as file:
            records.extend(json.loads(file.read()))
        return records

    @classmethod
    def _read_dev_records(cls) -> list:
        with open(join(cls.data_dir(), 'd-dev.json'), 'r', encoding='utf8') \
                as file:
            records = json.loads(file.read())
        with open(join(cls.data_dir(), 'm-dev.json'), 'r', encoding='utf8') \
                as file:
            records.extend(json.loads(file.read()))
        return records

    @classmethod
    def _read_test_records(cls) -> list:
        with open(join(cls.data_dir(), 'test.json'), 'r', encoding='utf8') \
                as file:
            records = json.loads(file.read())
        return records

    @classmethod
    def _make_prompt(cls, record, mask_target: bool) -> List[str]:
        res = []
        text = '\n'.join(record[0])
        questions = record[1]
        for each in questions:
            question = each['question']
            answer = cls.__mask_token if mask_target else each['answer']
            complete = cls.__make_single_prompt(text, question, answer)
            res.append(complete)
        return res

    def get_task_maker(self) -> TaskMaker:
        return PerplexityPromptShotsChoiceMakingTaskMaker(
            token_limit=context_limit,
            tokenizer=tokenizer,
            shots_generator=CommonShotsGenerator(
                length_sorted_example_tokens=self.get_sorted_example_tokens()
            ),
        )

    @classmethod
    def data_dir(cls) -> str:
        return abspath(join(__file__, '../..'))

    @classmethod
    def cache_dir(cls) -> str:
        return join(cls.work_dir(), 'cache')

    @classmethod
    def work_dir(cls) -> str:
        return abspath(join(__file__, '..'))

    @classmethod
    def __make_single_prompt(cls, text, question, answer) -> str:
        return f'阅读文本：“{text}”。问：{question}答：{answer}'
